// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <nebula/acero/exec_plan.h>
#include <nebula/acero/options.h>
#include <nebula/core/api.h>
#include <turbo/utility/status.h>
#include <nebula/compute/api.h>
#include <nebula/compute/expression.h>
#include <nebula/future/async_generator.h>
#include <nebula/future/future.h>

#include <cstdlib>
#include <iostream>
#include <memory>

// Demonstrate registering an Nebula compute function outside of the Nebula source tree

namespace cp = ::nebula::compute;
namespace ac = ::nebula::acero;

class ExampleFunctionOptionsType : public cp::FunctionOptionsType {
    const char *type_name() const override { return "ExampleFunctionOptionsType"; }

    std::string stringify(const cp::FunctionOptions &) const override {
        return "ExampleFunctionOptionsType";
    }

    bool compare(const cp::FunctionOptions &, const cp::FunctionOptions &) const override {
        return true;
    }

    std::unique_ptr<cp::FunctionOptions> copy(const cp::FunctionOptions &) const override;
    // optional: support for serialization
    // turbo::Result<std::shared_ptr<Buffer>> serialize(const FunctionOptions&) const override;
    // turbo::Result<std::unique_ptr<FunctionOptions>> deserialize(const Buffer&) const override;
};

cp::FunctionOptionsType *GetExampleFunctionOptionsType() {
    static ExampleFunctionOptionsType options_type;
    return &options_type;
}

class ExampleFunctionOptions : public cp::FunctionOptions {
public:
    ExampleFunctionOptions() : cp::FunctionOptions(GetExampleFunctionOptionsType()) {}
};

std::unique_ptr<cp::FunctionOptions> ExampleFunctionOptionsType::copy(
        const cp::FunctionOptions &) const {
    return std::make_unique<ExampleFunctionOptions>();
}

turbo::Status ExampleFunctionImpl(cp::KernelContext *ctx, const cp::ExecSpan &batch,
                                  cp::ExecResult *out) {
    out->value = batch[0].array.to_array_data();
    return turbo::OkStatus();
}

class ExampleNodeOptions : public ac::ExecNodeOptions {
};

// a basic ExecNode which ignores all input batches
class ExampleNode : public ac::ExecNode {
public:
    ExampleNode(ExecNode *input, const ExampleNodeOptions &)
            : ExecNode(/*plan=*/input->plan(), /*inputs=*/{input},
            /*input_labels=*/{"ignored"},
            /*output_schema=*/input->output_schema()) {}

    const char *kind_name() const override { return "ExampleNode"; }

    turbo::Status start_producing() override { return output_->input_finished(this, 0); }

    void resume_producing(ExecNode *output, int32_t counter) override {
        inputs_[0]->resume_producing(this, counter);
    }

    void pause_producing(ExecNode *output, int32_t counter) override {
        inputs_[0]->pause_producing(this, counter);
    }

    turbo::Status stop_producing_impl() override { return turbo::OkStatus(); }

    turbo::Status input_received(ExecNode *input, cp::ExecBatch batch) override {
        return turbo::OkStatus();
    }

    turbo::Status input_finished(ExecNode *input, int total_batches) override {
        return turbo::OkStatus();
    }
};

turbo::Result<ac::ExecNode *> example_exec_node_factory(ac::ExecPlan *plan,
                                                        std::vector<ac::ExecNode *> inputs,
                                                        const ac::ExecNodeOptions &options) {
    const auto &example_options =
            turbo::checked_cast<const ExampleNodeOptions &>(options);

    return plan->emplace_node<ExampleNode>(inputs[0], example_options);
}

const cp::FunctionDoc func_doc{
        "Example function to demonstrate registering an out-of-tree function",
        "",
        {"x"},
        "ExampleFunctionOptions"};

turbo::Status run_compute_register(int argc, char **argv) {
    const std::string name = "compute_register_example";
    auto func = std::make_shared<cp::ScalarFunction>(name, cp::Arity::Unary(), func_doc);
    cp::ScalarKernel kernel({nebula::int64()}, nebula::int64(), ExampleFunctionImpl);
    kernel.mem_allocation = cp::MemAllocation::NO_PREALLOCATE;
    TURBO_RETURN_NOT_OK(func->add_kernel(std::move(kernel)));

    auto registry = cp::get_function_registry();
    TURBO_RETURN_NOT_OK(registry->add_function(std::move(func)));

    nebula::Int64Builder builder(nebula::default_memory_pool());
    std::shared_ptr<nebula::Array> arr;
    TURBO_RETURN_NOT_OK(builder.append(42));
    TURBO_RETURN_NOT_OK(builder.finish(&arr));
    auto options = std::make_shared<ExampleFunctionOptions>();
    auto maybe_result = cp::call_function(name, {arr}, options.get());
    TURBO_RETURN_NOT_OK(maybe_result.status());

    std::cout << maybe_result->make_array()->to_string() << std::endl;

    // Expression serialization will raise NotImplemented if an expression includes
    // FunctionOptions for which serialization is not supported.
    auto expr = cp::call(name, {}, options);
    auto maybe_serialized = cp::serialize(expr);
    std::cerr << maybe_serialized.status().to_string() << std::endl;

    auto exec_registry = ac::default_exec_factory_registry();
    TURBO_RETURN_NOT_OK(
            exec_registry->add_factory("compute_register_example", example_exec_node_factory));

    auto maybe_plan = ac::ExecPlan::create();
    TURBO_RETURN_NOT_OK(maybe_plan.status());
    TURBO_MOVE_OR_RAISE(auto plan, maybe_plan);

    nebula::AsyncGenerator<std::optional<cp::ExecBatch>> source_gen, sink_gen;
    TURBO_RETURN_NOT_OK(
            ac::Declaration::sequence(
                    {
                            {"source",                   ac::SourceNodeOptions{nebula::schema({}), source_gen}},
                            {"compute_register_example", ExampleNodeOptions{}},
                            {"sink",                     ac::SinkNodeOptions{&sink_gen}},
                    })
                    .add_to_plan(plan.get())
                    .status());
    return turbo::OkStatus();
}

int main(int argc, char **argv) {
    auto status = run_compute_register(argc, argv);
    if (!status.ok()) {
        std::cerr << status.to_string() << std::endl;
        return EXIT_FAILURE;
    }
    return EXIT_SUCCESS;
}
